
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <time.h>
#include "grouping.h"

// 自定义的 strdup 实现
static char* hg_strdup(const char* s) {
    if (!s) return NULL;
    size_t len = strlen(s) + 1;
    char* new_str = malloc(len);
    if (new_str) {
        memcpy(new_str, s, len);
    }
    return new_str;
}

// 简单哈希函数
static uint32_t simple_hash(const char* str, uint32_t seed) {
    uint32_t hash = seed;
    while (*str) {
        hash = (hash * 31) + *str;
        str++;
    }
    return hash;
}

// 初始化HeavyGuardian
HeavyGuardian* hg_init(HGConfig config) {
    HeavyGuardian* hg = malloc(sizeof(HeavyGuardian));
    if (!hg) return NULL;
    
    hg->config = config;
    hg->total_count = 0;
    
    // 分配桶数组
    hg->buckets = malloc(config.depth * sizeof(HGBucket*));
    if (!hg->buckets) {
        free(hg);
        return NULL;
    }
    
    // 初始化随机种子
    srand(time(NULL));
    
    for (uint32_t d = 0; d < config.depth; d++) {
        hg->buckets[d] = malloc(config.width * sizeof(HGBucket));
        if (!hg->buckets[d]) {
            // 清理已分配的资源
            for (uint32_t i = 0; i < d; i++) {
                free(hg->buckets[i]);
            }
            free(hg->buckets);
            free(hg);
            return NULL;
        }
        
        for (uint32_t w = 0; w < config.width; w++) {
            // 分配高频项计数器
            hg->buckets[d][w].heavy = malloc(config.heavy_size * sizeof(HGCounter));
            if (!hg->buckets[d][w].heavy) {
                // 清理已分配的资源
                for (uint32_t i = 0; i <= d; i++) {
                    for (uint32_t j = 0; j < (i == d ? w : config.width); j++) {
                        free(hg->buckets[i][j].heavy);
                        free(hg->buckets[i][j].sparse);
                    }
                    free(hg->buckets[i]);
                }
                free(hg->buckets);
                free(hg);
                return NULL;
            }
            
            // 分配稀疏项计数器
            hg->buckets[d][w].sparse = malloc(config.sparse_size * sizeof(HGCounter));
            if (!hg->buckets[d][w].sparse) {
                // 清理已分配的资源
                for (uint32_t i = 0; i <= d; i++) {
                    for (uint32_t j = 0; j < (i == d ? w : config.width); j++) {
                        free(hg->buckets[i][j].heavy);
                        free(hg->buckets[i][j].sparse);
                    }
                    free(hg->buckets[i]);
                }
                free(hg->buckets);
                free(hg);
                return NULL;
            }
            
            // 初始化计数器
            for (uint32_t i = 0; i < config.heavy_size; i++) {
                hg->buckets[d][w].heavy[i].item = NULL;
                hg->buckets[d][w].heavy[i].count = 0;
            }
            
            for (uint32_t i = 0; i < config.sparse_size; i++) {
                hg->buckets[d][w].sparse[i].item = NULL;
                hg->buckets[d][w].sparse[i].count = 0;
            }
        }
    }
    
    return hg;
}

// 插入词项到HeavyGuardian
void hg_insert(HeavyGuardian* hg, const char* item) {
    hg->total_count++;
    
    for (uint32_t d = 0; d < hg->config.depth; d++) {
        // 计算哈希值确定桶位置
        uint32_t bucket_idx = simple_hash(item, d) % hg->config.width;
        HGBucket* bucket = &hg->buckets[d][bucket_idx];
        
        // 检查是否已在高频部分
        for (uint32_t i = 0; i < hg->config.heavy_size; i++) {
            if (bucket->heavy[i].item && strcmp(bucket->heavy[i].item, item) == 0) {
                bucket->heavy[i].count++;
                return;
            }
        }
        
        // 检查是否已在稀疏部分
        for (uint32_t i = 0; i < hg->config.sparse_size; i++) {
            if (bucket->sparse[i].item && strcmp(bucket->sparse[i].item, item) == 0) {
                bucket->sparse[i].count++;
                
                // 检查是否需要提升到高频部分
                uint32_t min_heavy_count = UINT32_MAX;
                uint32_t min_heavy_idx = 0;
                
                for (uint32_t j = 0; j < hg->config.heavy_size; j++) {
                    if (bucket->heavy[j].count < min_heavy_count) {
                        min_heavy_count = bucket->heavy[j].count;
                        min_heavy_idx = j;
                    }
                }
                
                if (bucket->sparse[i].count > min_heavy_count) {
                    // 交换稀疏项和高频项
                    HGCounter temp = bucket->heavy[min_heavy_idx];
                    bucket->heavy[min_heavy_idx] = bucket->sparse[i];
                    bucket->sparse[i] = temp;
                }
                
                return;
            }
        }
        
        // 不在任何部分，尝试插入稀疏部分
        // 找到稀疏部分中计数最小的项
        uint32_t min_sparse_count = UINT32_MAX;
        uint32_t min_sparse_idx = 0;
        bool has_empty = false;
        
        for (uint32_t i = 0; i < hg->config.sparse_size; i++) {
            if (bucket->sparse[i].item == NULL) {
                has_empty = true;
                min_sparse_idx = i;
                break;
            }
            
            if (bucket->sparse[i].count < min_sparse_count) {
                min_sparse_count = bucket->sparse[i].count;
                min_sparse_idx = i;
            }
        }
        
        if (has_empty || min_sparse_count == 0) {
            // 有空位或计数为0，直接插入
            if (bucket->sparse[min_sparse_idx].item) {
                free(bucket->sparse[min_sparse_idx].item);
            }
            char* new_item = hg_strdup(item);
            if (!new_item) {
                // 内存分配失败，跳过插入
                return;
            }
            bucket->sparse[min_sparse_idx].item = new_item;
            bucket->sparse[min_sparse_idx].count = 1;
        } else {
            // 应用指数衰减
            // 找到高频部分中最小的计数
            uint32_t min_heavy_count = UINT32_MAX;
            for (uint32_t i = 0; i < hg->config.heavy_size; i++) {
                if (bucket->heavy[i].count < min_heavy_count) {
                    min_heavy_count = bucket->heavy[i].count;
                }
            }
            
            // 计算衰减概率 P = base^(-min_heavy_count)
            double decay_prob = pow(hg->config.decay_base, -min_heavy_count);
            
            // 以衰减概率减少稀疏项计数
            if ((double)rand() / RAND_MAX < decay_prob) {
                if (bucket->sparse[min_sparse_idx].count > 0) {
                    bucket->sparse[min_sparse_idx].count--;
                }
                
                if (bucket->sparse[min_sparse_idx].count == 0) {
                    // 计数减为0，替换该项
                    free(bucket->sparse[min_sparse_idx].item);
                    char* new_item = hg_strdup(item);
                    if (!new_item) {
                        // 内存分配失败，跳过插入
                        bucket->sparse[min_sparse_idx].item = NULL;
                        return;
                    }
                    bucket->sparse[min_sparse_idx].item = new_item;
                    bucket->sparse[min_sparse_idx].count = 1;
                }
            }
        }
    }
}

// 查询词项频率
uint32_t hg_query(const HeavyGuardian* hg, const char* item) {
    uint32_t max_count = 0;
    for (uint32_t d = 0; d < hg->config.depth; d++) {
        uint32_t bucket_idx = simple_hash(item, d) % hg->config.width;
        const HGBucket* bucket = &hg->buckets[d][bucket_idx];
        
        // 检查高频部分
        for (uint32_t i = 0; i < hg->config.heavy_size; i++) {
            if (bucket->heavy[i].item && strcmp(bucket->heavy[i].item, item) == 0) {
                if (bucket->heavy[i].count > max_count) {
                    max_count = bucket->heavy[i].count;
                }
                break; // 在同一深度中找到后跳出内层循环
            }
        }
        
        // 检查稀疏部分
        for (uint32_t i = 0; i < hg->config.sparse_size; i++) {
            if (bucket->sparse[i].item && strcmp(bucket->sparse[i].item, item) == 0) {
                if (bucket->sparse[i].count > max_count) {
                    max_count = bucket->sparse[i].count;
                }
                break; // 在同一深度中找到后跳出内层循环
            }
        }
    }
    
    return max_count;
}

// 遍历HeavyGuardian中所有词项并调用回调函数
void hg_get_all_items(const HeavyGuardian* hg, void (*callback)(const char* item, uint32_t count)) {
    if (!hg || !callback) return;
    
    for (uint32_t d = 0; d < hg->config.depth; d++) {
        for (uint32_t w = 0; w < hg->config.width; w++) {
            const HGBucket* bucket = &hg->buckets[d][w];
            for (uint32_t i = 0; i < hg->config.heavy_size; i++) {
                if (bucket->heavy[i].item) {
                    callback(bucket->heavy[i].item, bucket->heavy[i].count);
                }
            }
            for (uint32_t i = 0; i < hg->config.sparse_size; i++) {
                if (bucket->sparse[i].item) {
                    callback(bucket->sparse[i].item, bucket->sparse[i].count);
                }
            }
        }
    }
}

// 释放HeavyGuardian资源
void hg_free(HeavyGuardian* hg) {
    if (!hg) return;
    
    for (uint32_t d = 0; d < hg->config.depth; d++) {
        for (uint32_t w = 0; w < hg->config.width; w++) {
            for (uint32_t i = 0; i < hg->config.heavy_size; i++) {
                free(hg->buckets[d][w].heavy[i].item);
            }
            free(hg->buckets[d][w].heavy);
            
            for (uint32_t i = 0; i < hg->config.sparse_size; i++) {
                free(hg->buckets[d][w].sparse[i].item);
            }
            free(hg->buckets[d][w].sparse);
        }
        free(hg->buckets[d]);
    }
    free(hg->buckets);
    free(hg);
}

// 获取HeavyGuardian的内存使用情况
size_t hg_memory_usage(const HeavyGuardian* hg) {
    if (!hg) return 0;
    
    size_t total = sizeof(HeavyGuardian);
    total += hg->config.depth * sizeof(HGBucket*);
    
    for (uint32_t d = 0; d < hg->config.depth; d++) {
        total += hg->config.width * sizeof(HGBucket);
        
        for (uint32_t w = 0; w < hg->config.width; w++) {
            total += hg->config.heavy_size * sizeof(HGCounter);
            total += hg->config.sparse_size * sizeof(HGCounter);
            
            // 计算字符串占用的内存
            for (uint32_t i = 0; i < hg->config.heavy_size; i++) {
                if (hg->buckets[d][w].heavy[i].item) {
                    total += strlen(hg->buckets[d][w].heavy[i].item) + 1;
                }
            }
            
            for (uint32_t i = 0; i < hg->config.sparse_size; i++) {
                if (hg->buckets[d][w].sparse[i].item) {
                    total += strlen(hg->buckets[d][w].sparse[i].item) + 1;
                }
            }
        }
    }
    
    return total;
}

// 重置HeavyGuardian（清空所有计数器）
void hg_reset(HeavyGuardian* hg) {
    if (!hg) return;
    
    hg->total_count = 0;
    
    for (uint32_t d = 0; d < hg->config.depth; d++) {
        for (uint32_t w = 0; w < hg->config.width; w++) {
            for (uint32_t i = 0; i < hg->config.heavy_size; i++) {
                free(hg->buckets[d][w].heavy[i].item);
                hg->buckets[d][w].heavy[i].item = NULL;
                hg->buckets[d][w].heavy[i].count = 0;
            }
            
            for (uint32_t i = 0; i < hg->config.sparse_size; i++) {
                free(hg->buckets[d][w].sparse[i].item);
                hg->buckets[d][w].sparse[i].item = NULL;
                hg->buckets[d][w].sparse[i].count = 0;
            }
        }
    }
}